#!/usr/bin/env python
import ssummolib
import sys
import Bio.SeqIO
from numpy import mean,std
import os.path

class Lengths( ):
    def __init__(self):
        self.lens = []
        self.mean = 0.
        self.nBases = 0
        self.minimum = 1E99
        self.maximum = 0
        self.name = u''
    def __add__(self,self2):
        end = Lengths()
        end.lens = self.lens + self2.lens
        end.mean = ( self.mean + self2.mean ) * 0.5
        end.nBases = self.nBases + self2.nBases
        end.minimum = min( self.minimum , self2.minimum )
        end.maximum = max( self.maximum , self2.maximum )
        return end
    def __str__(self):
        return '{0} {1} {2} {3} {4} {5} {6} = {7}'.format( self.name.ljust(28),str(self.minimum).ljust(6),str(self.maximum).ljust(7), str( len(self.lens)).ljust(8), '{0:.3f}'.format(self.getMean()).ljust(9) , '{0:.3f}'.format(self.getSTD()).ljust(10), self.nBases , bases(self.nBases) )
    def getLens(self, inhandle ,format='fasta'):
        for seq in Bio.SeqIO.parse( inhandle ,format ):
            break   # Just get the first sequence to check whether it's an alignment.
        if '-' in seq.seq.tostring():
            alignment = True  # If the first sequence has '-' characters, assume this is an alignment file, and we only want the actual residues.
        else:
            alignment = False
        inhandle.seek( 0 )
        if alignment:
            for seq in Bio.SeqIO.parse( inhandle, format ):
                seqLen = len( seq.seq.tostring().replace('-','') ) 
                self.lens.append( seqLen )
                self.minimum = min( seqLen , self.minimum )
                self.maximum = max( seqLen , self.maximum )
                self.nBases += seqLen
        else:
            for seq in Bio.SeqIO.parse( inhandle, format ):
                seqLen = len( seq )
                self.lens.append( seqLen )
                self.minimum = min( seqLen , self.minimum )
                self.maximum = max( seqLen , self.maximum )
                self.nBases += seqLen
        return self.lens
    def getMean( self ):
        if self.mean == 0.:
            self.mean = mean( self.lens )
        return self.mean
    def getSTD( self ):
        return std( self.lens )
    def nBases( self ):
        if self.nBases == 0:
            self.nBases = sum( self.nBases )
        return self.nBases


def get_stats( file_name,lengths = None ):
    lengths = Lengths()
    with file( file_name, 'r' ) as inFile:
        lens = lengths.getLens(inFile )
    lengths.name = os.path.realpath( inFile.name ).rsplit( os.path.sep , 1 )[ -1 ]
#    print 'lengths: ',lengths.getLens()
    nseqs = len( lens )
    mean = lengths.getMean()
    std = lengths.getSTD()
    return ( nseqs, mean, std )

def bases( nBases ):
    nBases = int( nBases )
    if nBases > 1E12: # Tb
        nBases = float( nBases / 1E12 )
        return '{0:.6f} Tb'.format(nBases)
    elif nBases > 1E9: # Gb
        return '{0:.6f} Gb'.format( float( nBases / 1E9 ) )
    elif nBases > 1E6: # Mb
        return '{0:.3f} Mb'.format( float( nBases / 1E6 ) )
    elif nBases > 1E3: # kb
        return '{0:.3f} kb'.format( float( nBases / 1E3 ) )
    else:
        return nBases, ' b'

def iterfiles( file_names,format='fasta' ):
    for file_name in file_names:
        lengths = Lengths( )
        with file( file_name , 'r' ) as inFile:
            lengths.getLens( inFile,format=format )
        lengths.name = os.path.realpath( file_name ).rsplit( os.path.sep , 1 )[-1]
        yield lengths

def get_headers():
    header = '{0} {1} {2} {3} {4} {5} {6}\n'.format( 'File'.ljust(28) ,'Min'.ljust(6),'Max'.ljust(7), 'nSeqs'.ljust(8),'mean'.ljust(9), 'Std.'.ljust(10), 'nBases'.ljust(20) )
    header += '{0} {1} {2} {3} {4} {5}'.format('name'.ljust(28), 'len'.ljust(6),'len'.ljust(7),  ' '.ljust(8),'length'.ljust(9), 'deviation'.ljust(10) )
    return header

options = ssummolib.Options()
options.multiargs = ['-in']
options.singleargs = ['-format']
options.options.update( {'-in' : [] ,
                         '-format' : 'fasta' } )


def main( file_names, format='fasta' ):
    total_lengths = Lengths()
    print get_headers()
    for lengths in iterfiles( file_names,format ): 
        print lengths
        total_lengths += lengths
    if len( file_names ) > 1:
        sys.stdout.write('\n##  TOTALS  ##\n')
        print total_lengths


if __name__ == '__main__':
    options.parseArgs(sys.argv[1:])
    file_names = options['-in']
    main( file_names, options['-format'] )
